// coroutine.c
#include "coroutine.h"
#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <stddef.h>
#include <string.h>
#include <stdint.h> 
#if __APPLE__ && __MACH__
    #include <sys/ucontext.h>
#else 
    #include <ucontext.h>
#endif

#define MIN_STACK_SIZE (128*1024)
#define MAX_STACK_SIZE (1024*1024)
#define DEFAULT_COROUTINE 128
#define MAIN_CO_ID 0

#define MIN(a, b) ((a) > (b) ? (b) : (a))
#define MAX(a, b) ((a) < (b) ? (b) : (a))

struct coroutine;

// 每个线程的调度器
typedef struct schedule {
    int stsize;             // 栈大小
    int nco;                // 当前有几个协程
    int cap;                // 协程数组容量
    int running;            // 当前正在运行的协程ID
    struct coroutine **co;  // 协程数组
} schedule_t;

// 协程数据
typedef struct coroutine {
    co_func func;           // 协程回调函数
    void *ud;               // 用户数据
    int pco;                // 前一个协程，即resume这个协程的那个协程
    ucontext_t ctx;         // 协程的执行环境
    schedule_t * sch;       // 调度器
    int status;             // 当前状态：CO_STATUS_RUNNING...
    char *stack;            // 栈内存
} coroutine_t;

schedule_t *co_open(int stsize) {
    schedule_t *S = malloc(sizeof(*S));
    S->nco = 0;
    S->stsize = MIN(MAX(stsize, MIN_STACK_SIZE), MAX_STACK_SIZE);
    S->cap = DEFAULT_COROUTINE;
    S->co = malloc(sizeof(coroutine_t *) * S->cap);
    memset(S->co, 0, sizeof(coroutine_t *) * S->cap);

    // 创建主协程
    int id = co_new(S, NULL, NULL);
    assert(id == MAIN_CO_ID);
    // 主协程为运行状态
    coroutine_t *co = S->co[MAIN_CO_ID];
    co->status = CO_STATUS_RUNNING;
    S->running = id;
    return S;
}

void co_close(schedule_t *S) {
    assert(S->running == MAIN_CO_ID);
    int i;
    for (i=0;i<S->cap;i++) {
        coroutine_t * co = S->co[i];
        if (co) {
            free(co->stack);
            free(co);
        }
    }
    free(S->co);
    S->co = NULL;
    free(S);
}

static void cofunc(uint32_t low32, uint32_t hi32) {
    uintptr_t ptr = (uintptr_t)low32 | ((uintptr_t)hi32 << 32);
    schedule_t *S = (schedule_t *)ptr;
    int id = S->running;
    coroutine_t *co = S->co[id];
    co->func(S, co->ud);
    // 标记协程为死亡
    co->status = CO_STATUS_DEAD;
    --S->nco;
    // 恢复前一个协程
    coroutine_t *pco = S->co[co->pco];
    pco->status = CO_STATUS_RUNNING;
    S->running = co->pco;
    ucontext_t dummy;
    swapcontext(&dummy, &pco->ctx);
}

int co_new(schedule_t *S, co_func func, void *ud) {
    int cid = -1;
    if (S->nco >= S->cap) {
        cid = S->cap;
        S->co = realloc(S->co, S->cap * 2 * sizeof(coroutine_t *));
        memset(S->co + S->cap , 0 , sizeof(coroutine_t *) * S->cap);
        S->cap *= 2;
    } else {
        int i;
        for (i=0;i<S->cap;i++) {
            int id = (i+S->nco) % S->cap;
            if (S->co[id] == NULL) {
                cid = id;
                break;
            } 
            else if (S->co[id]->status == CO_STATUS_DEAD) {
                // printf("reuse dead coroutine: %d\n", id);
                cid = id;
                break;
            }
        }
    }

    if (cid >= 0) {
        coroutine_t *co;
        if (S->co[cid])
            co = S->co[cid];
        else {
            co = malloc(sizeof(*co));
            co->pco = 0;
            co->stack = cid != MAIN_CO_ID ? malloc(S->stsize) : 0;
            S->co[cid] = co;
        }
        ++S->nco;

        co->func = func;
        co->ud = ud;
        co->sch = S;
        co->status = CO_STATUS_SUSPEND;

        if (func) {
            coroutine_t *curco = S->co[S->running];
            assert(curco);
            getcontext(&co->ctx);
            co->ctx.uc_stack.ss_sp = co->stack;
            co->ctx.uc_stack.ss_size = S->stsize;
            co->ctx.uc_link = &curco->ctx;
            uintptr_t ptr = (uintptr_t)S;
            makecontext(&co->ctx, (void (*)(void))cofunc, 2, (uint32_t)ptr, (uint32_t)(ptr>>32));
        }
    }

    return cid;
}

int co_resume(schedule_t * S, int id) {
    assert(id >=0 && id < S->cap);
    coroutine_t *co = S->co[id];
    coroutine_t *curco = S->co[S->running];
    if (co == NULL || curco == NULL)
        return -1;
    int status = co->status;
    switch(status) {
    case CO_STATUS_SUSPEND:
        curco->status = CO_STATUS_NORMAL;
        co->pco = S->running;
        co->status = CO_STATUS_RUNNING;
        S->running = id;
        swapcontext(&curco->ctx, &co->ctx);
        return 0;
    default:
        return -1;
    }
}

int co_yield(schedule_t * S) {
    int id = S->running;
    // 主协程不能yield
    if (id == MAIN_CO_ID)
        return -1;
    // 恢复当前协程环境
    assert(id >= 0);
    coroutine_t * co = S->co[id];
    coroutine_t *pco = S->co[co->pco];
    co->status = CO_STATUS_SUSPEND;
    pco->status = CO_STATUS_RUNNING;
    S->running = co->pco;
    swapcontext(&co->ctx ,&pco->ctx);
    return 0;
}

int co_status(schedule_t * S, int id) {
    assert(id>=0 && id < S->cap);
    if (S->co[id] == NULL) {
        return CO_STATUS_DEAD;
    }
    return S->co[id]->status;
}

int co_running(schedule_t * S) {
    return S->running;
}
